import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split


# 创建数据集：实心圆（0）和方形圆（1）
def generate_circle_image(shape, size=28):
    """生成实心圆或方形圆的图像"""
    image = torch.ones((size, size), dtype=torch.float32)  # 默认背景为1（白色）

    if shape == 'circle':
        # 画实心圆
        y, x = torch.meshgrid(torch.arange(size), torch.arange(size), indexing='ij')  # 显式指定 indexing 参数
        center = (size // 2, size // 2)
        radius = size // 4
        mask = (x - center[1]) ** 2 + (y - center[0]) ** 2 <= radius ** 2
        image[mask] = 0  # 圆的部分为0（黑色）

    elif shape == 'square':
        # 画方形圆（即方形内切圆）
        square_size = size // 2
        top_left = (size // 4, size // 4)
        bottom_right = (top_left[0] + square_size, top_left[1] + square_size)
        image[top_left[0]:bottom_right[0], top_left[1]:bottom_right[1]] = 0  # 方形部分为0（黑色）

    return image


def create_dataset(num_samples=1000):
    images = []
    labels = []

    for _ in range(num_samples // 2):
        # 生成实心圆
        images.append(generate_circle_image('circle'))
        labels.append(0)  # 实心圆 -> 类别 0

        # 生成方形圆
        images.append(generate_circle_image('square'))
        labels.append(1)  # 方形圆 -> 类别 1

    images = torch.stack(images)  # 转为张量
    labels = torch.tensor(labels, dtype=torch.long)  # 标签转为长整型张量

    # 打乱数据
    indices = torch.randperm(len(images))
    images = images[indices]
    labels = labels[indices]

    # 将数据拆分为训练集和测试集
    train_size = int(0.8 * len(images))
    X_train, X_test = images[:train_size], images[train_size:]
    y_train, y_test = labels[:train_size], labels[train_size:]

    # 添加通道维度
    X_train = X_train.unsqueeze(1)
    X_test = X_test.unsqueeze(1)

    return X_train, y_train, X_test, y_test


# 定义卷积神经网络
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 2)  # 二分类输出层

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))  # 卷积 + 激活 + 池化
        x = self.pool(torch.relu(self.conv2(x)))  # 卷积 + 激活 + 池化
        x = x.view(-1, 64 * 7 * 7)  # 展平
        x = torch.relu(self.fc1(x))  # 全连接层1 + 激活
        x = self.fc2(x)  # 全连接层2（输出）
        return x


# 创建数据集
X_train, y_train, X_test, y_test = create_dataset()

# 初始化模型
model = SimpleCNN()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 10
for epoch in range(num_epochs):
    model.train()

    # 前向传播
    outputs = model(X_train)

    # 计算损失
    loss = criterion(outputs, y_train)

    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 输出每个epoch的损失
    if (epoch + 1) % 2 == 0:
        print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}")

# 评估模型
model.eval()
with torch.no_grad():
    outputs = model(X_test)
    _, predicted = torch.max(outputs, 1)
    accuracy = (predicted == y_test).sum().item() / y_test.size(0)
    print(f"Test Accuracy: {accuracy * 100:.2f}%")

# 可视化一些测试样本
# 可视化一些测试样本
fig, axes = plt.subplots(1, 5, figsize=(12, 3))
for i in range(5):
    # 将张量转换为 numpy 数组并使用 matplotlib 显示
    axes[i].imshow(X_test[i].cpu().numpy().squeeze(), cmap='gray')
    axes[i].set_title(f"Pred: {predicted[i].item()} - True: {y_test[i].item()}")
    axes[i].axis('off')
plt.show()

